flydsl: gfx942 FP8 MQA logits indexer kernel (+ Triton FN/FNUZ fix) by akii96 · Pull Request #3913 · ROCm/aiter

akii96 · 2026-06-25T02:34:41Z

Motivation

DeepSeek-V4 Flash's sparse attention uses a lightning indexer (fp8 MQA logits). This PR adds a FlyDSL gfx942 drop-in for the Triton fp8_mqa_logits that is faster on the DSv4 prefill shapes for both the mono-dtype config (q/k both e4m3fnuz) as well as the currently vllm upstream's mixed config.

The uplift does not depend on operand dtype: even with both operands fnuz (no conversion on either side), FlyDSL is 1.5–2.5× faster at the kernel level and cuts end-to-end prefill TTFT by up to 17% at 128K and increases with increasing ctx size.

Technical Details

New FlyDSL gfx942 kernel which is a drop-in for the Triton one (same args, return, and clean_logits semantics).
(Optional safety net) Triton fp8_mqa_logits FN/FNUZ fix by @vpietila-amd; a no-op for fnuz/fnuz, so existing callers are unchanged. This only matters while callers still send mixed q=fn/k=fnuz (see note below); it can be dropped if maintainers prefer the source-side fix only.

FlyDSL kernel initially translated from GEAK by @peyron-amd, with perf iterated by myself and @vpietila-amd.

Test Plan

pytest op_tests/flydsl_tests/test_flydsl_fp8_mqa_logits.py — validated against the torch reference and cross-checked vs the Triton kernel (calc_diff < 1e-3): both fp8 dtype combos

Test Result

335/335 pass.

Kernel speedup (gfx942, H=64, D=128). FlyDSL wins in both operand configs:

Mixed q=e4m3fn / k=e4m3fnuz (current vLLM behavior):

S_q × S_kv	Triton	FlyDSL	speedup
4096 × 4096	0.47 ms	0.21 ms	2.2×
8192 × 8192	1.48 ms	0.70 ms	2.1×
128 × 32768	0.30 ms	0.12 ms	2.4×
671 × 131072	3.14 ms	1.94 ms	1.6×

Mono fnuz/fnuz (the production config; gfx942, H=64, D=128):

S_q × S_kv	Triton	FlyDSL	speedup
4096 × 4096	0.34 ms	0.22 ms	1.5×
8192 × 8192	1.15 ms	0.74 ms	1.6×
128 × 32768	0.29 ms	0.12 ms	2.5×
671 × 131072	3.14 ms	1.78 ms	1.8×

End-to-end DeepSeek-V4 Flash TTFT (TP4, MI300-class gfx942, OSL=27, concurrency=1; Triton gfx942 indexer vs FlyDSL, leveraging PR 46730 on vLLM which fixes the mixed dtype):

ISL	Triton	FlyDSL	speedup	TTFT reduction
32K	2099 ms	1981 ms	1.06×	−5.7%
128K	11337 ms	9401 ms	1.21×	−17.0%

The E2E gain grows with context length (the indexer logits cost scales with N), which is where prefill latency hurts most.

Note on the mixed FN/FNUZ path (re: @valarLip)

You were right to push on the mixed dtype usage. The mixed q=fn/k=fnuz inputs come from vLLM, not aiter. Currently DSv4-Flash's fused_indexer_q.py quantizes Q to e4m3fn while the gfx942 KV cache is fnuz. The proper fix is at the source: vllm-project/vllm#46730 makes Q fnuz so both operands match. The Triton FN/FNUZ fix here is an optional bridge for the current vLLM behavior and can be dropped if you'd rather rely on the source fix.

github-actions · 2026-06-25T02:35:29Z

🏷️ CI Guide

Runs automatically on every PR:

✅ Pre-checks (submodule verification, code formatting)
✅ Aiter op tests (gfx942 + gfx950)
✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
`ci:triton-300x`	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
`ci:sglang`	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
`ci:atom`	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
`ci:atom_full`	ATOM accuracy suite for PR and main models from ATOM `models_accuracy.json`
`ci:vllm`	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
`ci:all`	All standard extended tests (excludes `ci:atom_full`)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 3913 --add-label <label>

valarLip · 2026-06-25T05:14:31Z

why there is a mixed FN/FNUZ path?

akii96 · 2026-06-25T07:19:57Z

@valarLip

The mixed types come from vllm. On gfx942 the KV cache is fnuz, but DSv4 Flash quantizes Q to fn (fused_indexer_q.py only emits fn), so the kernel gets q=fn / k=fnuz.

We can fix that at the source and give the Q-quant an fnuz path like the KV cache already has (I can look at it later today!)

Even with both fnuz (no conversion on either side), flydsl is faster (compared below) so my thinking is that we can still have this PR

fnuz/fnuz	triton	flydsl
4096×4096	0.34 ms	0.22 ms (1.5×)
8192×8192	1.15 ms	0.74 ms (1.6×)
128×32768	0.29 ms	0.12 ms (2.5×)
671×131072	3.14 ms	1.78 ms (1.8×)

I am still polishing this work up but do you think the flydsl kernel improvements for the mono dtype are upstream worthy or do you have other concerns ? 😅

akii96 · 2026-06-26T01:06:44Z

@valarLip I've updated the PR description significantly to give a fuller picture of this work, could you have a second look and advise how you'd like to scope this PR?

flydsl kernel only: drop the Triton FN/FNUZ fix and rely on the vLLM source-side fix (Q→fnuz) to eliminate the mixed dtype, or
flydsl kernel + keep the mixed-dtype bridge: the earlier source-side attempts upstream never merged and went stale, and the current one ([ROCm][Perf][Bugfix] DSv4 indexer: use platform FP8 dtype (fnuz) for Q-quant on gfx942 vllm-project/vllm#46730) isn't guaranteed to land, so carrying mixed-dtype support in aiter short-term may be a necessary evil

Thankful for any insight you can provide here!

vpietila-amd · 2026-06-26T06:54:53Z

+        convert_kv_fn = KV.dtype != _fnuz
+        scale_mul = 1.0
+        if convert_q_fn:
+            scale_mul *= 2.0


We should add

if arch == "gfx942":

such that we don't accidentally convert FP8-FN on other architectures that support FP8.

fixed now in 99bd4d3! Thanks

vpietila-amd · 2026-06-26T07:08:53Z

+        _fnuz,
+        _fn,
+    ), f"Q/KV must be e4m3 fp8 (fnuz or fn); got {Q.dtype}, {KV.dtype}"
+    convert_q_fn = Q.dtype != _fnuz


We need to check

if arch == "gfx942":

to ensure that we don't convert to FP8-FNUZ on other architectures.

This kernel is optimized for gfx942 and while it should run also on other FP8 supported architectures, the performance will probably be suboptimal as the it is not using e.g. async loads that are available in new architectures. At minimum, we should test that the current implementation works on gfx950.

fixed now in 99bd4d3! Thanks. The decision to claify for others is to gate this kernel for now on gfx942

FlyDSL gfx942 kernel for the DeepSeek lightning indexer (fp8 MQA logits), a drop-in for the Triton fp8_mqa_logits. Includes an adaptive grid.y KV-column split for block-starved small-M/long-context prefill, in-kernel FP8 FN->FNUZ handling, and correctness tests vs the torch and Triton references. Also fixes the Triton fp8_mqa_logits mixed FN/FNUZ path (no-op for fnuz/fnuz inputs). Co-authored-by: Ville Pietilä <Ville.Pietila@amd.com> Co-authored-by: Kristoffer Peyron <Kristoffer.Peyron@amd.com> Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

Co-authored-by: Ville Pietilä <Ville.Pietila@amd.com> Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

vpietila-amd · 2026-06-26T08:59:28Z

+    WPB = waves_per_block
+    MR_BLOCK_THREADS = 64 * WPB
+
+    assert H % 16 == 0, f"num_heads={H} must be a multiple of 16 for MFMA"


We should add here named constants

MFMA_M = 16 MFMA_N = 16 MFMA_K = 32

And use those constants in the rest of the code. Now there too many magic number 16 and 32 floating around.

Yeah made sense! and now addressed in d0d58f0. Added MFMA_M/MFMA_N/MFMA_K and used them for the tile dims everywhere, and renamed the lane_div_16/lane_mod_16 vars to lane_div_N/lane_mod_N to match

Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

vpietila-amd

Reviewed the implementation and the new FlyDSL kernel looks good to me.

akii96 · 2026-06-30T07:26:38Z

@valarLip we gathered more results with the new FlyDSL kernel introduced in this PR

Outside of DSv4 Flash we also tried replacing the Triton kernel with the FlyDSL counterpart for GLM-5.2-FP8:

Max Conc	TTFT (ms)
2	20389.85 → 15427.41 (−24%)
4	32335.73 → 25395.11 (−21%)
6	33781.63 → 27584.69 (−18%)
8	39559.43 → 30936.41 (−22%)
GeoMean	30637.75 → 24046.14 (−21%)

TTFT reduction is consistent (18–24%) across concurrency levels

Hoping to get your review on this and get this merged as we noticed similar improvements across DS3.2/GLM5.1 as well (happy to provide more results on request!)

valarLip · 2026-07-07T08:10:14Z

please refact ut follow this one https://github.com/ROCm/aiter/tree/main/.claude/skills/aiter-op-test

Refactored in 9e8797b!

Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

…3913) * flydsl: add gfx942 fp8 MQA logits indexer kernel FlyDSL gfx942 kernel for the DeepSeek lightning indexer (fp8 MQA logits), a drop-in for the Triton fp8_mqa_logits. Includes an adaptive grid.y KV-column split for block-starved small-M/long-context prefill, in-kernel FP8 FN->FNUZ handling, and correctness tests vs the torch and Triton references. Also fixes the Triton fp8_mqa_logits mixed FN/FNUZ path (no-op for fnuz/fnuz inputs). Co-authored-by: Ville Pietilä <Ville.Pietila@amd.com> Co-authored-by: Kristoffer Peyron <Kristoffer.Peyron@amd.com> Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> * flydsl: fix i32 index overflow in fp8 MQA logits output write Co-authored-by: Ville Pietilä <Ville.Pietila@amd.com> Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> * flydsl/triton: gate FN->FNUZ conversion to gfx942 Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> * flydsl: name MFMA tile-dim constants (MFMA_M/N/K) Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> * flydsl: refactor fp8_mqa_logits test to aiter-op-test standard Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> --------- Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> Co-authored-by: Ville Pietilä <Ville.Pietila@amd.com>

… (#49544) Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

… (vllm-project#49544) Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

… (vllm-project#49544) Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> Signed-off-by: Wyett <wyettzeng@gmail.com>

… (vllm-project#49544) Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com>

… (vllm-project#49544) Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> Signed-off-by: khushali9 <khushali.desai9@gmail.com>

akii96 requested review from aghamari and vpietila-amd June 25, 2026 02:35

akii96 force-pushed the aakif/fp8-mqa-logits branch 2 times, most recently from 976b50f to 138d402 Compare June 25, 2026 23:17

akii96 marked this pull request as ready for review June 26, 2026 00:53

akii96 requested a review from a team June 26, 2026 00:53

akii96 requested a review from coderfeli June 26, 2026 01:08

vpietila-amd reviewed Jun 26, 2026

View reviewed changes

akii96 and others added 3 commits June 26, 2026 11:08

flydsl: fix i32 index overflow in fp8 MQA logits output write

19537a3

Co-authored-by: Ville Pietilä <Ville.Pietila@amd.com> Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

flydsl/triton: gate FN->FNUZ conversion to gfx942

079edff

Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

akii96 force-pushed the aakif/fp8-mqa-logits branch from 99bd4d3 to 079edff Compare June 26, 2026 08:08

vpietila-amd reviewed Jun 26, 2026

View reviewed changes

flydsl: name MFMA tile-dim constants (MFMA_M/N/K)

d0d58f0

Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

vpietila-amd previously approved these changes Jun 26, 2026

View reviewed changes

valarLip requested changes Jul 7, 2026

View reviewed changes

akii96 dismissed vpietila-amd’s stale review via 3ad3e13 July 7, 2026 10:56

flydsl: refactor fp8_mqa_logits test to aiter-op-test standard

9e8797b

Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

akii96 force-pushed the aakif/fp8-mqa-logits branch from 3ad3e13 to 9e8797b Compare July 7, 2026 11:01

valarLip approved these changes Jul 7, 2026

View reviewed changes

valarLip merged commit 0571aea into main Jul 7, 2026
56 checks passed

valarLip deleted the aakif/fp8-mqa-logits branch July 7, 2026 13:49

fhuizing mentioned this pull request Jul 13, 2026

[FlyDSL] Paged mla indexer #4221

Open

akii96 mentioned this pull request Jul 23, 2026

[ROCm][Perf] gfx942: use FlyDSL fp8 MQA logits kernel (ROCm/aiter#3913) vllm-project/vllm#49544

Merged

yraparti mentioned this pull request Aug 3, 2026

rocke implementation of fp8 mqa logits on gfx942 ROCm/rocm-libraries#10287

Draft

1 task

vpietila-amd mentioned this pull request Aug 3, 2026

[FlyDSL] gfx950 FP8 MQA logits indexer kernel #4538

Open

1 task

tjtanaa pushed a commit to vllm-project/vllm that referenced this pull request Aug 16, 2026

[ROCm][Perf] gfx942: use FlyDSL fp8 MQA logits kernel (ROCm/aiter#3913)…

6914d60

… (#49544) Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

zyp2014 pushed a commit to zyp2014/vllm that referenced this pull request Aug 21, 2026

[ROCm][Perf] gfx942: use FlyDSL fp8 MQA logits kernel (ROCm/aiter#3913)…

64c6ec5

… (vllm-project#49544) Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

wyettzeng pushed a commit to wyettzeng/vllm that referenced this pull request Aug 21, 2026

[ROCm][Perf] gfx942: use FlyDSL fp8 MQA logits kernel (ROCm/aiter#3913)…

60e9b1f

… (vllm-project#49544) Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> Signed-off-by: Wyett <wyettzeng@gmail.com>

Uh oh!

Conversation

akii96 commented Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Note on the mixed FN/FNUZ path (re: @valarLip)

Uh oh!

github-actions Bot commented Jun 25, 2026

🏷️ CI Guide

Uh oh!

valarLip commented Jun 25, 2026

Uh oh!

akii96 commented Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

akii96 commented Jun 26, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vpietila-amd left a comment

Choose a reason for hiding this comment

Uh oh!

akii96 commented Jun 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

akii96 commented Jun 25, 2026 •

edited

Loading

akii96 commented Jun 25, 2026 •

edited

Loading

akii96 commented Jun 30, 2026 •

edited

Loading